import pandas as pd
import numpy as np
import statsmodels.api as sm
import matplotlib.pyplot as plt
import seaborn as sns
from arch import arch_model
import shap
from sklearn.metrics import mean_absolute_error

plt.style.use('ggplot')
sns.set(rc={'font.sans-serif': 'SimHei'})


class QuantAnalyzer:
    def __init__(self, data_path):
        self.df = pd.read_csv(data_path, parse_dates=['trade_date'])
        self.y = self.df['price_ret']
        self.X = sm.add_constant(self.df[['flow_lag_1', 'rsi', 'index_ret']])

    def _diagnose_model(self, model):
        """模型诊断三件套"""
        residuals = model.resid

        # 正态性检验
        jb_test = sm.stats.jarque_bera(residuals)
        print(f"Jarque-Bera正态性检验: p={jb_test:.3f}")

        # 异方差检验
        white_test = sm.stats.diagnostic.het_white(residuals, self.X)
        print(f"White异方差检验: p={white_test:.3f}")

        # 波动聚集检验
        garch = arch_model(residuals * 100, vol='Garch', p=1, q=1).fit(disp='off')
        print(f"GARCH系数: α={garch.params['alpha']:.3f} (p={garch.pvalues['alpha']:.3f})")

        # 可视化诊断
        fig, ax = plt.subplots(1, 2, figsize=(12, 4))
        sm.qqplot(residuals, line='s', ax=ax)
        pd.plotting.autocorrelation_plot(residuals, ax=ax)
        plt.tight_layout()

    def _rolling_backtest(self, window=60):
        """滚动窗口回测"""
        preds, actuals = [], []
        for i in range(window, len(self.df)):
            train_X = self.X.iloc[i - window:i]
            train_y = self.y.iloc[i - window:i]

            model = sm.OLS(train_y, train_X).fit()
            pred = model.predict(self.X.iloc[[i]])

            preds.append(pred)
            actuals.append(self.y.iloc[i])

        mae = mean_absolute_error(actuals, preds)
        print(f"滚动预测MAE: {mae:.4f}")

        # 可视化预测效果
        pd.DataFrame({'实际': actuals, '预测': preds}).plot(
            title="滚动预测效果对比", figsize=(10, 5)
        )

    def _shap_analysis(self, model):
        """特征贡献度分析"""
        explainer = shap.Explainer(model)
        shap_values = explainer(self.X)

        plt.figure(figsize=(10, 4))
        shap.summary_plot(shap_values, self.X, plot_type="bar")
        plt.title("SHAP特征重要性")

    def analyze(self):
        """分析主流程"""
        # 建模
        model = sm.OLS(self.y, self.X).fit()
        print(model.summary())

        # 诊断
        self._diagnose_model(model)

        # 滚动回测
        self._rolling_backtest()

        # SHAP分析
        self._shap_analysis(model)

        plt.show()


if __name__ == "__main__":
    analyzer = QuantAnalyzer("cleaned_data.csv")
    analyzer.analyze()
